#ifndef VOLUME_RENDERER_H
#define VOLUME_RENDERER_H

#include <G3D/G3DAll.h>
#include <GLG3D/GLG3D.h>
#include <fstream>

struct ColorPoint
{
public:
	bool sel;
	float x;
	Color3uint8 y;

public:
	ColorPoint() : sel(false), x(0.0f), y(Color3())
	{
	}

	ColorPoint(float _x, const Color3uint8 &_y) : sel(false), x(_x), y(_y)
	{
	}
};

class TFEditor : public GuiControl
{
public:
	float layoutT, layoutB;
	int mode, x0, y0;
	Color3uint8 pick;
	Array<ColorPoint> points;
	Array<float> histogram;
	GImage function;
	GImage hslspace;
	TextureRef colormap;
	TextureRef spacemap;

public:
	bool save();
	bool load();
	void insert(float x, const Color3uint8 &y);
	void update();
	template<typename T>
	void stat(const Array<T> &data, T minVal, T maxVal, float scale);
	virtual void render(RenderDevice *rd, const GuiThemeRef &skin) const;
	virtual bool onEvent(const GEvent &event);

public:
	TFEditor(GuiContainer *parent = NULL, const GuiText &text = "", int width = 512, int height = 256);
};

class TFWindow
{
public:
	GuiWindowRef window;
	TFEditor *editor;

public:
	TFWindow()
	{
		window = GuiWindow::create("Transfer Function", 0, Rect2D::xywh(0, 0, 256, 256), GuiTheme::NORMAL_WINDOW_STYLE, GuiWindow::HIDE_ON_CLOSE);
		GuiPane *pane = window->pane();
		editor = new TFEditor(pane);
		pane->addCustom(editor);
		pane->addNumberBox("Red:", &editor->pick.r, "", GuiTheme::LINEAR_SLIDER, uint8(0), uint8(255))->setWidth(editor->rect().width());
		pane->addNumberBox("Green:", &editor->pick.g, "", GuiTheme::LINEAR_SLIDER, uint8(0), uint8(255))->setWidth(editor->rect().width());
		pane->addNumberBox("Blue:", &editor->pick.b, "", GuiTheme::LINEAR_SLIDER, uint8(0), uint8(255))->setWidth(editor->rect().width());
		pane->pack();
		pane->window()->pack();
	}
};

class VolumeRenderer
{
public:
	float Ka, Kd, Ks, Kp;
	float alphaScale;
	float colorScale;
	float valueScale;
	float valueShift;
	float sampleStep;
	ShaderRef shader;

public:
	void render(RenderDevice *rd, const Rect2D &rect, TextureRef volume, TextureRef transfer, const Vector3 &boxMin, const Vector3 &boxMax)
	{
		shader->args.set("Aspect", rect.width() / rect.height());
		shader->args.set("CotFOV", rd->projectionMatrix().row(1).y);
		shader->args.set("InvMat", rd->modelViewMatrix().inverse().toMatrix4());
		Matrix4 p = rd->projectionMatrix();
		Matrix4 w = rd->objectToWorldMatrix();
		Matrix4 v = rd->cameraToWorldMatrix();
		shader->args.set("AlphaScale", alphaScale);
		shader->args.set("ColorScale", colorScale);
		shader->args.set("ValueScale", valueScale);
		shader->args.set("ValueShift", valueShift);
		shader->args.set("SampleStep", sampleStep);
		shader->args.set("Ka", Ka);
		shader->args.set("Kd", Kd);
		shader->args.set("Ks", Ks);
		shader->args.set("Kp", Kp);
		shader->args.set("BoxMin", boxMin);
		shader->args.set("BoxMax", boxMax);
		shader->args.set("VoxelSize", Vector3(1.0f / volume->width(), 1.0f / volume->height(), 1.0f / volume->depth()));
		shader->args.set("Projection", rd->projectionMatrix());
		shader->args.set("Volume", volume);
		shader->args.set("Transfer", transfer);
		rd->pushState();
		rd->setBlendFunc(RenderDevice::BLEND_SRC_ALPHA, RenderDevice::BLEND_ONE_MINUS_SRC_ALPHA);
		//rd->setBlendFunc(RenderDevice::BLEND_ONE, RenderDevice::BLEND_ONE_MINUS_SRC_ALPHA);
		rd->setShader(shader);
		rd->push2D();
		rd->setDepthTest(RenderDevice::DEPTH_LEQUAL);
		rd->setDepthWrite(true);
		Draw::fastRect2D(rect, rd);
		rd->pop2D();
		rd->popState();
	}

public:
	VolumeRenderer()
	{
		alphaScale = 1.0f;
		colorScale = 1.0f;
		valueScale = 1.0f;
		valueShift = 0.0f;
		sampleStep = 0.005f;
		Ka = 0.2f;
		Kd = 1.0f;
		Ks = 1.0f;
		Kp = 100.0f;
		shader = Shader::fromFiles("VolumeRenderer.vp", "VolumeRenderer.fp");
		assert(shader->ok());
	}
};

class Viewer : public GuiControl
{
public:
	int mode;
	short x0, y0;
	CoordinateFrame frame;
	TextureRef color;
	TextureRef depth;

public:
	virtual void render(RenderDevice *rd, const GuiThemeRef &skin) const
	{
		skin->pauseRendering();
		rd->setTexture(0, color);
		Draw::fastRect2D(m_rect, rd);
		rd->setTexture(0, NULL);
		skin->resumeRendering();
	}

	virtual bool onEvent(const GEvent &event)
	{
		if(event.type == GEventType::MOUSE_BUTTON_DOWN)
		{
			mode = event.button.button + 1;
			x0 = (short)event.button.x;
			y0 = (short)event.button.y;
			return true;
		}

		if(event.type == GEventType::MOUSE_BUTTON_UP)
		{
			mode = 0;
			return true;
		}

		if((mode == 1) && (event.type == GEventType::MOUSE_MOTION))
		{
			int x1 = (short)event.button.x;
			int y1 = (short)event.button.y;
			float dx = (x1 - x0) * 0.5f;
			float dy = (y1 - y0) * 0.5f;
			frame.rotation = Matrix3::fromAxisAngle(Vector3::unitY(), toRadians(dx)) * frame.rotation;
			frame.rotation = Matrix3::fromAxisAngle(Vector3::unitX(), toRadians(dy)) * frame.rotation;
			frame.rotation.orthonormalize();
			x0 = x1;
			y0 = y1;
			return true;
		}

		if((mode == 2) && (event.type == GEventType::MOUSE_MOTION))
		{
			int x1 = (short)event.button.x;
			int y1 = (short)event.button.y;
			float dx = (x1 - x0) * 0.01f;
			float dy = (y1 - y0) * 0.01f;
			frame.translation.x += dx;
			frame.translation.y -= dy;
			x0 = x1;
			y0 = y1;
			return true;
		}

		if((mode == 3) && (event.type == GEventType::MOUSE_MOTION))
		{
			int x1 = (short)event.button.x;
			int y1 = (short)event.button.y;
			float dx = (x1 - x0) * 0.05f;
			float dy = (y1 - y0) * 0.05f;
			frame.translation.z -= dy;
			if(frame.translation.z > 0.0f)
				frame.translation.z = 0.0f;
			x0 = x1;
			y0 = y1;
			return true;
		}

		return false;
	}

public:
	Viewer(GuiContainer *parent = NULL, const GuiText &text = "", int width = 640, int height = 480) : GuiControl(parent, text)
	{
		setWidth(width);
		setHeight(height);
		color = Texture::createEmpty("", m_rect.width(), m_rect.height(), ImageFormat::RGBA8(), Texture::DIM_2D_NPOT, Texture::Settings::video());
		depth = Texture::createEmpty("", m_rect.width(), m_rect.height(), ImageFormat::DEPTH32(), Texture::DIM_2D_NPOT, Texture::Settings::shadow());
	}
};

class SliceWindow
{
public:
	GuiWindowRef window;
	Viewer *viewer;
	FramebufferRef fbo;
	ShaderRef shader;
	int sliceAxis;
	float sliceDepth;
	float colorScale;

public:
	void renderSlice(RenderDevice *rd, TextureRef volume, float valueShift = 0.0f, float valueScale = 1.0f)
	{
		fbo->set(Framebuffer::COLOR_ATTACHMENT0, viewer->color);
		fbo->set(Framebuffer::DEPTH_ATTACHMENT, viewer->depth);
		rd->pushState(fbo);
		rd->setCameraToWorldMatrix(CoordinateFrame());
		rd->setObjectToWorldMatrix(CoordinateFrame());
		rd->setProjectionMatrix(Matrix4::orthogonalProjection(0.0f, 1.0f, 0.0f, 1.0f, 1.0f, -1.0f));
		rd->setColorClearValue(Color4(1.0f, 1.0f, 1.0f, 1.0f));
		rd->setDepthTest(RenderDevice::DEPTH_ALWAYS_PASS);
		rd->setDepthWrite(false);
		rd->clear(true, true, false);
		shader->args.set("ValueShift", valueShift);
		shader->args.set("ValueScale", valueScale * colorScale);
		shader->args.set("Volume", volume);
		rd->setShader(shader);
		rd->beginPrimitive(RenderDevice::Primitive::QUADS);
		switch(sliceAxis)
		{
		case 0:
			rd->setTexCoord(0, Vector3(sliceDepth, 0.0f, 0.0f));
			rd->sendVertex(Vector2(0.0f, 0.0f));
			rd->setTexCoord(0, Vector3(sliceDepth, 1.0f, 0.0f));
			rd->sendVertex(Vector2(1.0f, 0.0f));
			rd->setTexCoord(0, Vector3(sliceDepth, 1.0f, 1.0f));
			rd->sendVertex(Vector2(1.0f, 1.0f));
			rd->setTexCoord(0, Vector3(sliceDepth, 0.0f, 1.0f));
			rd->sendVertex(Vector2(0.0f, 1.0f));
			break;
		case 1:
			rd->setTexCoord(0, Vector3(0.0f, sliceDepth, 0.0f));
			rd->sendVertex(Vector2(0.0f, 0.0f));
			rd->setTexCoord(0, Vector3(1.0f, sliceDepth, 0.0f));
			rd->sendVertex(Vector2(1.0f, 0.0f));
			rd->setTexCoord(0, Vector3(1.0f, sliceDepth, 1.0f));
			rd->sendVertex(Vector2(1.0f, 1.0f));
			rd->setTexCoord(0, Vector3(0.0f, sliceDepth, 1.0f));
			rd->sendVertex(Vector2(0.0f, 1.0f));
			break;
		case 2:
			rd->setTexCoord(0, Vector3(0.0f, 0.0f, sliceDepth));
			rd->sendVertex(Vector2(0.0f, 0.0f));
			rd->setTexCoord(0, Vector3(1.0f, 0.0f, sliceDepth));
			rd->sendVertex(Vector2(1.0f, 0.0f));
			rd->setTexCoord(0, Vector3(1.0f, 1.0f, sliceDepth));
			rd->sendVertex(Vector2(1.0f, 1.0f));
			rd->setTexCoord(0, Vector3(0.0f, 1.0f, sliceDepth));
			rd->sendVertex(Vector2(0.0f, 1.0f));
			break;
		}
		rd->endPrimitive();
		rd->popState();
	}

public:
	SliceWindow()
	{
		sliceAxis = 0;
		sliceDepth = 0.5f;
		colorScale = 1.0f;
		window = GuiWindow::create("Slice Renderer", 0, Rect2D::xywh(0, 0, 256, 256), GuiTheme::NORMAL_WINDOW_STYLE, GuiWindow::HIDE_ON_CLOSE);
		GuiPane *pane = window->pane();
		viewer = new Viewer(pane);
		pane->addCustom(viewer);
		Array<std::string> axisList("X-Axis", "Y-Axis", "Z-Axis");
		pane->addDropDownList("Slice Axis", axisList, &sliceAxis);
		pane->addNumberBox("Depth", &sliceDepth, "", GuiTheme::LINEAR_SLIDER, 0.0f, 1.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Scale", &colorScale, "", GuiTheme::LOG_SLIDER, 0.0f, 100.0f)->setWidth(viewer->rect().width());
		pane->pack();
		pane->window()->pack();
		fbo = Framebuffer::create("");
		shader = Shader::fromFiles("VolumeSlicer.vp", "VolumeSlicer.fp");
		assert(shader->ok());
	}
};

class VRWindow
{
public:
	GuiWindowRef window;
	Viewer *viewer;
	VolumeRenderer renderer;
	FramebufferRef fbo;
	SliceWindow sliceWindow;
	bool showBox, showSlice;

public:
	void renderBegin(RenderDevice *rd, TextureRef volume, TextureRef transfer, const Vector3 &boxMin, const Vector3 &boxMax)
	{
		fbo->set(Framebuffer::COLOR_ATTACHMENT0, viewer->color);
		fbo->set(Framebuffer::DEPTH_ATTACHMENT, viewer->depth);
		rd->pushState(fbo);
		////rd->setCameraToWorldMatrix(CoordinateFrame(Vector3(0.0f, 0.0f, (boxMax - boxMin).length())));
		//rd->setCameraToWorldMatrix(CoordinateFrame(Vector3(0.0f, 0.0f, 0.0f)));
		//rd->setObjectToWorldMatrix(CoordinateFrame(viewer->frame));
		rd->setCameraToWorldMatrix(viewer->frame.inverse());
		rd->setObjectToWorldMatrix(CoordinateFrame());
		float h = 0.01f;
		float w = h * viewer->rect().width() / viewer->rect().height();
		rd->setProjectionMatrix(Matrix4::perspectiveProjection(-w, w, -h, h, h * 2.0f, 100.0f));
		rd->setColorClearValue(Color4(0.0f, 0.0f, 0.0f, 1.0f));
		rd->setDepthClearValue(1.0f);
		rd->setDepthTest(RenderDevice::DEPTH_LEQUAL);
		rd->setDepthWrite(true);
		rd->clear(true, true, false);
		if(showBox)
		{
			glDisable(GL_LINE_SMOOTH);
			Draw::box(Box(boxMin, boxMax), rd, Color4::zero(), Color3::white());
			glEnable(GL_LINE_SMOOTH);
		}
		if(showSlice)
		{
			float sd;
			rd->setColor(Color3::gray());
			rd->setCullFace(RenderDevice::CULL_NONE);
			rd->beginPrimitive(RenderDevice::Primitive::QUADS);
			switch(sliceWindow.sliceAxis)
			{
			case 0:
				sd = boxMin.x + (boxMax.x - boxMin.x) * sliceWindow.sliceDepth;
				rd->sendVertex(Vector3(sd, boxMin.y, boxMin.z));
				rd->sendVertex(Vector3(sd, boxMax.y, boxMin.z));
				rd->sendVertex(Vector3(sd, boxMax.y, boxMax.z));
				rd->sendVertex(Vector3(sd, boxMin.y, boxMax.z));
				break;
			case 1:
				sd = boxMin.y + (boxMax.y - boxMin.y) * sliceWindow.sliceDepth;
				rd->sendVertex(Vector3(boxMin.x, sd, boxMin.z));
				rd->sendVertex(Vector3(boxMax.x, sd, boxMin.z));
				rd->sendVertex(Vector3(boxMax.x, sd, boxMax.z));
				rd->sendVertex(Vector3(boxMin.x, sd, boxMax.z));
				break;
			case 2:
				sd = boxMin.z + (boxMax.z - boxMin.z) * sliceWindow.sliceDepth;
				rd->sendVertex(Vector3(boxMin.x, boxMin.y, sd));
				rd->sendVertex(Vector3(boxMax.x, boxMin.y, sd));
				rd->sendVertex(Vector3(boxMax.x, boxMax.y, sd));
				rd->sendVertex(Vector3(boxMin.x, boxMax.y, sd));
				break;
			}
			rd->endPrimitive();
		}
	}

	void renderEnd(RenderDevice *rd, TextureRef volume, TextureRef transfer, const Vector3 &boxMin, const Vector3 &boxMax)
	{
		Rect2D rect = Rect2D::xywh(0, 0, viewer->rect().width(), viewer->rect().height());
		rd->setBlendFunc(RenderDevice::BLEND_SRC_ALPHA, RenderDevice::BLEND_ONE_MINUS_SRC_ALPHA);
		float sampleStep = renderer.sampleStep;
		if(viewer->mode != 0)
			renderer.sampleStep = 0.005f;
		renderer.render(rd, rect, volume, transfer, boxMin, boxMax);
		renderer.sampleStep = sampleStep;
		//Draw::axes(CoordinateFrame(boxMin), rd);
		rd->popState();
		sliceWindow.renderSlice(rd, volume, renderer.valueShift, renderer.valueScale);
	}

	void render(RenderDevice *rd, TextureRef volume, TextureRef transfer, const Vector3 &boxMin, const Vector3 &boxMax)
	{
		renderBegin(rd, volume, transfer, boxMin, boxMax);
		renderEnd(rd, volume, transfer, boxMin, boxMax);
	}

public:
	VRWindow()
	{
		window = GuiWindow::create("Volume Renderer", 0, Rect2D::xywh(0, 0, 256, 256), GuiTheme::NORMAL_WINDOW_STYLE, GuiWindow::HIDE_ON_CLOSE);
		GuiPane *pane = window->pane();
		viewer = new Viewer(pane);
		pane->addCustom(viewer);
		pane->addNumberBox("Alpha Scale", &renderer.alphaScale, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Color Scale", &renderer.colorScale, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Value Scale", &renderer.valueScale, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Value Shift", &renderer.valueShift, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Sample Step", &renderer.sampleStep, "", GuiTheme::LOG_SLIDER, 0.0001f, 0.1f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Ambient", &renderer.Ka, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Diffuse", &renderer.Kd, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Specular", &renderer.Ks, "", GuiTheme::LOG_SLIDER, 0.0f, 10.0f)->setWidth(viewer->rect().width());
		pane->addNumberBox("Shininess", &renderer.Kp, "", GuiTheme::LOG_SLIDER, 1.0f, 1000.0f)->setWidth(viewer->rect().width());
		pane->addCheckBox("Show Bounding Box", &showBox);
		sliceWindow.window->pane()->addCheckBox("Show Slice", &showSlice);
		pane->pack();
		pane->window()->pack();
		fbo = Framebuffer::create("");
		showBox = true;
		showSlice = false;
	}
};

#endif
